import cv2
import numpy as np
import os
from utils.basic import cv_imread, cv_imwrite

MIN_MATCH_COUNT = 4


def get_homography(img_1, img_2):
    """
    计算两幅图片之间的单应矩阵
    :return: 单应矩阵
    """
    sift = cv2.xfeatures2d.SIFT_create()
    kp1, des1 = sift.detectAndCompute(img_1, None)
    kp2, des2 = sift.detectAndCompute(img_2, None)

    FLANN_INDEX_KDTREE = 0
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)

    flann = cv2.FlannBasedMatcher(index_params, search_params)

    matches = flann.knnMatch(des1, des2, k=2)

    good = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good.append(m)

    if len(good) > MIN_MATCH_COUNT:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1,
                                                                         2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1,
                                                                         2)

        homo, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

        return homo
    else:
        raise RuntimeError("Not enough matches are found - %d/%d" % (
            len(good), MIN_MATCH_COUNT))


def __calc_corners(homo: np.array, img_shape: tuple):
    """
    计算变换后图像的四个顶点坐标
    :param homo: 变换矩阵
    :param img_shape: 原图大小
    :return: 变换后的坐标
    """
    ret = []
    # 变换后图片左上角坐标计算
    left_up_corner = np.array([0, 0, 1])
    left_up_corner = np.dot(homo, left_up_corner)
    ret.append([left_up_corner[0] / left_up_corner[2],
                left_up_corner[1] / left_up_corner[2]])

    # 变换后图片左下角坐标计算
    left_bottom_corner = np.array([0, img_shape[0], 1])
    left_bottom_corner = np.dot(homo, left_bottom_corner)
    ret.append([left_bottom_corner[0] / left_bottom_corner[2],
                left_bottom_corner[1] / left_bottom_corner[2]])

    # 变换后图片右上角坐标计算
    right_up_corner = np.array([img_shape[1], 0, 1], dtype=np.float64)
    right_up_corner = np.dot(homo, right_up_corner)
    ret.append([right_up_corner[0] / right_up_corner[2],
                right_up_corner[1] / right_up_corner[2]])

    # 变换后图片右下角坐标计算
    right_bottom_corner = np.array([img_shape[1], img_shape[0], 1],
                                   dtype=np.float64)
    right_bottom_corner = np.dot(homo, right_bottom_corner)
    ret.append([right_bottom_corner[0] / right_bottom_corner[2],
                right_bottom_corner[1] / right_bottom_corner[2]])

    return ret


def stitch_imgs(img_paths, output_path):
    """
    对输入的图片进行拼接
    :param img_paths: 所有要拼接的图片的路径
    :param output_path: 输出的结果图片的路径
    """
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    imgs = [cv_imread(name) for name in img_paths]

    # 对图片进行缩放
    if imgs[0].shape[0] >= 1000:
        multiple = imgs[0].shape[0] // 500
        for i in range(len(imgs)):
            imgs[i] = cv2.resize(imgs[i], (int(imgs[i].shape[1] / multiple),
                                           int(imgs[i].shape[0] / multiple)))

    img_num = len(imgs)
    center = img_num // 2
    img_gray = imgs

    # 求所有图片到相邻图片的单应矩阵
    homo_neig = []
    for i in range(img_num):
        if i < center:
            homo_neig.append(
                np.mat(get_homography(img_gray[i], img_gray[i + 1])))
        elif i > center:
            homo_neig.append(
                np.mat(get_homography(img_gray[i], img_gray[i - 1])))
        else:
            homo_neig.append(
                np.mat(get_homography(img_gray[i], img_gray[center])))

    # 求所有图片到中心图片的单应矩阵
    homo_cen = []
    for i in range(img_num):
        if i < center:
            h = homo_neig[center - 1]
            for j in range(center - 2, i - 1, -1):
                h = h * homo_neig[j]
        elif i > center:
            h = homo_neig[center + 1]
            for j in range(center + 2, i + 1):
                h = h * homo_neig[j]
        else:
            h = homo_neig[center]
        homo_cen.append(h)

    # 计算平移及缩放矩阵
    up_boundary, right_boundary, down_boundary, left_boundary = 0, 0, 0, 0
    for i in range(img_num):  # 求上边界
        corner = __calc_corners(np.array(homo_cen[i]),
                                (imgs[i].shape[:2]))
        up_boundary = min(up_boundary, int(corner[0][1]), int(corner[2][1]))
        right_boundary = max(right_boundary, int(corner[2][0]),
                             int(corner[3][0]))
        down_boundary = max(down_boundary, int(corner[1][1]),
                            int(corner[3][1]))
        left_boundary = min(left_boundary, int(corner[0][0]),
                            int(corner[1][0]))
    Tx = np.array(
        [[1.0, 0, -left_boundary], [0, 1.0, -up_boundary], [0, 0, 1.0]])
    for i in range(img_num):  # 将平移变换加入透视变换中
        homo_cen[i] = np.dot(Tx, homo_cen[i])

    # 将所有图片对中心图片做透视变换
    height = down_boundary - up_boundary
    width = right_boundary - left_boundary
    imgn2cen = []
    for i in range(img_num):
        imgn2cen.append(
            cv2.warpPerspective(imgs[i], homo_cen[i], (width, height)))

    # 所有图片投影到背景图片上
    dst = imgn2cen[0]
    for i in range(img_num):
        for row in range(height):
            for col in range(width):
                temp = dst[row, col]
                if temp[0] == 0 and temp[1] == 0 and temp[2] == 0:
                    dst[row, col] = imgn2cen[i][row, col]

    cv_imwrite(output_path, dst)


if __name__ == '__main__':
    imgs_name0 = ['0-1.jpg', '0-2.jpg', '0-3.jpg', '0-4.jpg', '0-5.jpg',
                  '0-6.jpg', '0-7.jpg', '0-8.jpg', '0-9.jpg']

    output_path = "D:/shelf detect/stitch/code/n_stitch/output.jpg"

    stitch_imgs(imgs_name0, output_path)
